import threading
import time

class threadPool:

    """Flexible thread pool class.  Creates a pool of threads, then
    accepts tasks that will be dispatched to the next available
    thread."""
    
    threadPoolSleepTime = 0.1

    def __init__(self, threadNum):

        """Initialize the thread pool with numThreads workers."""

        self.__threads = []
        self.__resizeLock = threading.Lock()
        self.__taskLock = threading.Lock()
        self.__tasks = []
        self.__isJoining = False
        self.__threadsNum = threadNum
        self.setThreadCount(threadNum)

    def setThreadCount(self, newThreadNum):

        """ External method to set the current pool size.  Acquires
        the resizing lock, then calls the internal version to do real
        work."""
        
        # Can't change the thread count if we're shutting down the pool!
        if self.__isJoining:
            return False;

        self.__resizeLock.acquire()
        try:
            self.__setThreadCountNolock(newThreadNum)
        finally:
            self.__resizeLock.release()
        
        return True

    def __setThreadCountNolock(self, newThreadNum):

        """Set the current pool size, spawning or terminating threads
        if necessary.  Internal use only; assumes the resizing lock is
        held."""
        
        while newThreadNum > len(self.__threads):
            newThread = threadPoolThread(self)
            self.__threads.append(newThread)
            newThread.start()

        while newThreadNum < len(self.__threads):
            # thread terminates after current cmd
            self.__threads[0].goAway()
            del self.__threads[0]

    def getThreadCount(self):

        """Return the number of threads in the pool."""

        self.__resizeLock.acquire()
        try:
            return len(self.__threads)
        finally:
            self.__resizeLock.release()

    def queueTask(self, task, args = None, taskCallBack = None):

        """Insert a task into the queue.  task must be callable;
        args and taskCallback can be None."""
        
        if self.__isJoining:
            return False
        if not callable(task):
            return False

        self.__taskLock.acquire()
        try:
            self.__tasks.append((task, args, taskCallBack))
            return True
        finally:
            self.__taskLock.release()

    def getNextTask(self):

        """ Retrieve the next task from the task queue.  For use
        only by ThreadPoolThread objects contained in the pool."""
        
        self.__taskLock.acquire()
        try:
            if self.__tasks == []:
                return (None, None, None)
            else:
                return self.__tasks.pop(0)
        finally:
            self.__taskLock.release()

    def joinAll(self, waitForTasks = True, waitForThreads = True):

        """ Clear the task queue and terminate all pooled threads,
        optionally allowing the tasks and threads to finish."""
        
        self.__isJoining = True

        if waitForTasks:
            while self.__tasks != []:
                time.sleep(threadPool.threadPoolSleepTime)

        self.__resizeLock.acquire()
        try:
            # __setThreadCountNolock(0) will let all threads finish the current cmd and then terminate. So it's always waitForThreads.
            self.__setThreadCountNolock(0)
            # unnecessary
            #if waitForThreads:
            #    for t in self.__threads:
            #        t.join()
            #        del t

            # Reset the pool for potential reuse
            self.__isJoining = False
        finally:
            self.__resizeLock.release()

class threadPoolThread(threading.Thread):

    """ Pooled thread class. """

    threadSleepTime = 0.1

    def __init__(self, pool):

        """ Initialize the thread and remember the pool. """

        threading.Thread.__init__(self)
        self.__pool = pool
        self.__isDying = False

    def run(self):

        """ Until told to quit, retrieve the next task and execute
        it, calling the callback if any.  """
        
        while not self.__isDying:
            cmd, args, callback = self.__pool.getNextTask()
            if cmd is None:
                # If there's nothing to do, just sleep a bit
                time.sleep(threadPoolThread.threadSleepTime)
            elif callback is None:
                cmd(args)
            else:
                callback(cmd(args))

    def goAway(self):

        """ Exit the run loop next time through."""

        self.__isDying = True

# Usage example
if __name__ == "__main__":

    from random import randrange

    def sortTask(data):
        print "sortTask starting for", data
        numbers = range(data[0], data[1])
        for a in numbers:
            rnd = randrange(0, len(numbers)-1)
            a, numbers[rnd] = numbers[rnd], a
        numbers.sort()
        print "sortTask done for", data
        return "sorter for ", data

    def waitTask(data):
        print "waitTask starting for", data
        time.sleep(data)
        print "waitTask done for", data
        return "waiter for ", data

    def taskCallback(data):
        print "taskCallback on ", data

    pool = threadPool(3)
    pool.queueTask(sortTask, (111, 1111), taskCallback)
    pool.queueTask(waitTask, 3)
    pool.joinAll()

    print "reuse pool"
    pool.setThreadCount(4)
    pool.queueTask(sortTask, (1, 3457))
    pool.queueTask(waitTask, 1)
    pool.joinAll()
